# coding:utf-8

import os
import random
import argparse

#rootdir = 'E:/A_dy_bishe/PV_heartbeat_detection/mydataset/'
parser = argparse.ArgumentParser()
#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
#parser.add_argument('--xml_path', default='E:/A_dy_bishe/PV_heartbeat_detection/mydataset/Annotations', type=str, help='input xml label path')
#数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='D:\\Data\\20250514_SpecParam_EA\\labelme_to_yolo', type=str, help='output txt label path')
#感觉这个参数没什么用，但是我懒得改了
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.8
#xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
img_path="D:\\Data\\20250514_SpecParam_EA\\labelme_to_yolo"
#新数据加进去记得把图放进去
total_img = os.listdir(img_path)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_img)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
file_path_train=open(r"D:\\Data\\20250514_SpecParam_EA\\labelme_to_yolo\\train.txt","w")
file_path_val=open(r"D:\\Data\\20250514_SpecParam_EA\\labelme_to_yolo\\val.txt","w")
for i in list_index:
    name = total_img[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)

        if i in train:
            file_train.write(name)
            file_path_train.write(img_path+"/"+total_img[i])
            file_path_train.write("\n")
        else:
            file_val.write(name)
            file_path_val.write(img_path+"/"+total_img[i])
            file_path_val.write("\n")
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
